from typing import Dict, List, Optional, Union
import pandas as pd
import os


class Node:
    """
    表示图中的一个节点，包含从数据文件中对应的所有字段。
    """
    
    def __init__(
        self,
        id: int,
        op: str,
        buf_id: Optional[int] = None,
        size: Optional[Union[int, float]] = None,
        type: Optional[str] = None,
        pipe: Optional[str] = None,
        cycles: Optional[int] = None,
        bufs: Optional[List[int]] = None
    ):
        self.id: int = id
        self.op: str = op
        self.buf_id: Optional[int] = buf_id
        self.size: Optional[Union[int, float]] = size
        self.type: Optional[str] = type
        self.pipe: Optional[str] = pipe
        self.cycles: Optional[int] = cycles
        self.bufs: List[int] = bufs if bufs is not None else []
    
    def __repr__(self) -> str:
        return f"Node(id={self.id}, op='{self.op}', buf_id={self.buf_id}, type='{self.type}')"
    
    def __str__(self) -> str:
        return f"Node {self.id}: {self.op}"


class Graph:
    """
    表示一个有向图，包含节点和边的信息。
    """
    
    def __init__(self):
        # 存储所有节点，以节点ID为键
        self.nodes: Dict[int, Node] = {}
        
        # 邻接表：存储每个节点的后继节点列表
        self.successors: Dict[int, List[int]] = {}
        
        # 邻接表：存储每个节点的前驱节点列表  
        self.predecessors: Dict[int, List[int]] = {}
    
    def add_node(self, node: Node) -> None:
        """
        添加一个节点到图中。
        
        Args:
            node: 要添加的Node对象
        """
        self.nodes[node.id] = node
        
        # 初始化邻接表中的条目
        if node.id not in self.successors:
            self.successors[node.id] = []
        if node.id not in self.predecessors:
            self.predecessors[node.id] = []
    
    def add_edge(self, source_id: int, dest_id: int) -> None:
        """
        添加一条从source_id到dest_id的有向边。
        
        Args:
            source_id: 源节点的ID
            dest_id: 目标节点的ID
        """
        # 确保节点存在于邻接表中
        if source_id not in self.successors:
            self.successors[source_id] = []
        if dest_id not in self.predecessors:
            self.predecessors[dest_id] = []
        if source_id not in self.predecessors:
            self.predecessors[source_id] = []
        if dest_id not in self.successors:
            self.successors[dest_id] = []
        
        # 添加边（如果不存在的话）
        if dest_id not in self.successors[source_id]:
            self.successors[source_id].append(dest_id)
        if source_id not in self.predecessors[dest_id]:
            self.predecessors[dest_id].append(source_id)
    
    def get_node(self, node_id: int) -> Optional[Node]:
        """
        根据ID获取节点。
        
        Args:
            node_id: 节点ID
            
        Returns:
            对应的Node对象，如果不存在则返回None
        """
        return self.nodes.get(node_id)
    
    def get_successors(self, node_id: int) -> List[int]:
        """
        获取指定节点的所有后继节点ID。
        
        Args:
            node_id: 节点ID
            
        Returns:
            后继节点ID的列表
        """
        return self.successors.get(node_id, [])
    
    def get_predecessors(self, node_id: int) -> List[int]:
        """
        获取指定节点的所有前驱节点ID。
        
        Args:
            node_id: 节点ID
            
        Returns:
            前驱节点ID的列表
        """
        return self.predecessors.get(node_id, [])
    
    def __len__(self) -> int:
        """返回图中节点的数量。"""
        return len(self.nodes)
    
    def __repr__(self) -> str:
        return f"Graph(nodes={len(self.nodes)}, edges={sum(len(succ) for succ in self.successors.values())})"
    
    def __str__(self) -> str:
        return f"Graph with {len(self.nodes)} nodes and {sum(len(succ) for succ in self.successors.values())} edges"


def load_graph_from_csv(task_name: str, data_dir: str = 'data/CSV-version') -> Graph:
    """
    从CSV文件中读取数据并创建一个填充好的Graph对象。
    
    Args:
        task_name: 任务名称（如 "Conv_Case0"）
        data_dir: 数据目录路径
        
    Returns:
        填充了节点和边数据的Graph对象
    """
    # 构建文件路径
    nodes_file = os.path.join(data_dir, f"{task_name}_Nodes.csv")
    edges_file = os.path.join(data_dir, f"{task_name}_Edges.csv")
    
    # 检查文件是否存在
    if not os.path.exists(nodes_file):
        raise FileNotFoundError(f"Nodes文件不存在: {nodes_file}")
    if not os.path.exists(edges_file):
        raise FileNotFoundError(f"Edges文件不存在: {edges_file}")
    
    # 创建Graph实例
    graph = Graph()
    
    # 读取Nodes CSV文件
    nodes_df = pd.read_csv(nodes_file)
    
    # 处理节点数据
    for _, row in nodes_df.iterrows():
        # 处理bufs字段 - 从字符串解析为整数列表
        bufs = []
        if pd.notna(row.get('Bufs')) and str(row['Bufs']).strip():
            bufs_str = str(row['Bufs']).strip().strip('"')
            if bufs_str:
                # 解析逗号分隔的数字字符串
                try:
                    bufs = [int(x.strip()) for x in bufs_str.split(',') if x.strip()]
                except ValueError:
                    # 如果解析失败，保持空列表
                    bufs = []
        
        # 创建Node对象，处理空值
        node = Node(
            id=int(row['Id']),
            op=str(row['Op']),
            buf_id=int(row['BufId']) if pd.notna(row.get('BufId')) else None,
            size=float(row['Size']) if pd.notna(row.get('Size')) else None,
            type=str(row['Type']) if pd.notna(row.get('Type')) else None,
            pipe=str(row['Pipe']) if pd.notna(row.get('Pipe')) else None,
            cycles=int(row['Cycles']) if pd.notna(row.get('Cycles')) else None,
            bufs=bufs
        )
        
        # 添加节点到图中
        graph.add_node(node)
    
    # 读取Edges CSV文件
    edges_df = pd.read_csv(edges_file)
    
    # 处理边数据
    for _, row in edges_df.iterrows():
        start_node_id = int(row['StartNodeId'])
        end_node_id = int(row['EndNodeId'])
        
        # 添加边到图中
        graph.add_edge(start_node_id, end_node_id)
    
    return graph
